from mcp.client.streamable_http import streamablehttp_client
from mcp.client.session import ClientSession
from mcp.client.auth import OAuthClientProvider
from mcp.shared.auth import OAuthClientMetadata
from storage import MemoryTokenStorage
import asyncio
import webbrowser
from loguru import logger
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from threading import Thread
from queue import SimpleQueue
from pydantic import AnyUrl


mcp_server_url = "http://localhost:8080"
auth_server_url = "http://localhost:9090"

queue : SimpleQueue = SimpleQueue()

class AuthServerHTTPRequestHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        # print(self.path)
        parsed = urlparse(self.path)
        query = parse_qs(parsed.query)
        if "code" in query:
            code = query.get("code")[0]
            state = query.get("state")[0]
            queue.put({"code": code, "state": state})
            print("code:{}, state:{}".format(code, state))
        self.send_response(200)
        self.send_header('Content-type', 'text/html')
        self.end_headers()
        self.wfile.write(b"Authorization successful.")


# 授权服务器会返回一个/authorize?参数，客户端要访问这个url
async def redirect_handler(authorization_url: str) -> None:
    # 这个函数的作用，就是想办法访问/authorize?参数，这个url。
    #一定要通过浏览器去访问。

    # 在浏览器打开授权页面进行登陆时，就开启3030端口的服务，以防止授权完后，客户端的3030服务还未开启
    server = HTTPServer(('localhost', 3030), AuthServerHTTPRequestHandler)
    thread = Thread(target=server.serve_forever, daemon=True, name="callback_server")
    thread.start()

    logger.info(f"Redirecting to {authorization_url}")
    webbrowser.open(authorization_url)


async def callback_handler() -> tuple[str|None, str|None]:
    # 这个函数中，要想办法获取到：http://localhost:3030/callback?code=xx&state上的code参数和state参数。
    # http://localhost:3030：这个路由是由我们客户端自己实现的服务。授权服务器在验证用户名和密码通过后会访问这个url
    info = queue.get(timeout=60*30) # get方法，在队列没有值的时候，会阻塞线程，一直等直到有数据
    if "code" in info and "state" in info:
        return info["code"], info["state"]


async def main():

    # oauth_provider
    oauth_provider = OAuthClientProvider(
        server_url=mcp_server_url,
        client_metadata=OAuthClientMetadata(
            redirect_uris=[AnyUrl("http://localhost:3030/callback")],
            scope="user",
        ),
        storage=MemoryTokenStorage(),
        redirect_handler=redirect_handler,
        callback_handler=callback_handler,
    )

    async with streamablehttp_client(
        url=f"{mcp_server_url}/mcp",
        auth=oauth_provider,
    ) as (read_stream, write_stream, get_session_id):

        async with ClientSession(read_stream, write_stream) as session:
            # 对session进行初始化
            await session.initialize()
            # 请求 url（f"{mcp_server_url}/mcp"）
            tools = await session.list_tools()
            print(tools)


if __name__ == '__main__':
    asyncio.run(main())